#! /usr/bin/perl -w
#
# Probe for checking GridFTP server via file transfer
# Copyright (c) 2006 Emir Imamagic
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

#
# Changes and Modifications
# =========================
# 22-Aug-2006 - Created;
#
# 26-Mar-2007 - Common libraries; plugin is WLCG-compatible
#
# 29-May-2007 - Updates related to WLCG specification changes
#             - -u,--url parameters replaced with --path
#
# 19-Sep-2007 - Allow for {glite,edg}-gridftp-rm commands variation
#               to support glite 3.1 UI
#
# 11-Jan-2008 - Added wrapper for safe execution of shell commands
#
# 15-Sep-2008 - changed sgutils namespace
#
# 12-Mar-2009 - switched from globus-url-copy and edg-gridftp-rm to uberftp.
#
# 16-Dec-2009 - 1.7 Migrated to Apache 2.0 license
#
# 23-Aug-2010 - Migrated to Nagios::Plugin.

use strict;
use Nagios::Plugin;
use Sys::Hostname;
use GridMon::sgutils qw($VONAME &checkProxy &processCommand);
use Sys::Syslog;
use TOM::Nagios qw(nagios_debug);

local $SIG{__DIE__}  = \&TOM::Nagios::handle_die;
local $SIG{__WARN__} = \&TOM::Nagios::handle_warn;
openlog("GridFTP-probe", "ndelay,pid", "user");
nagios_debug("started");

# Standard variables used in Nagios::Plugin constructor
use constant PROGNAME => "GridFTP-probe";
use constant VERSION => '1.8';
use constant DESCRIPTION => 'Probe for checking GridFTP server via file transfer';
use constant EXTRA_DESC => "";
use constant LICENSE => 'This nagios plugin is free software, and comes with ABSOLUTELY NO WARRANTY.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
   http://www.apache.org/licenses/LICENSE-2.0
Copyright 2009 Emir Imamagic';
use constant SHORTNAME => 'GridFTP';
use constant USAGE => "usage: $0 [ -v ] \n";
use constant DEFAULT_PORT => 2811;
use constant DEFAULT_TMP_DIR => '/tmp';

my %COMMANDS;

sub processErrorOutput {
    my $res = shift;
    my $add = shift || "";
    my $answer = "$add failed with error: ";

    $res =~ s/^\s*\n//mg;
    $res =~ s/^error: //mg;
    $res =~ s/\n//mg;

    $answer .= $res;

    $answer;
}

sub createFile {
    my $localurl = shift;
    my $data = shift;
    my $answer;
    my $state = OK;

    unless (open (STD,"> $localurl")) {
        $answer = "Creating local test file failed.\n";
        $state = CRITICAL;
    } else {
        print STD $data;
        unless (close (STD)) {
            $answer = "Closing local test file failed.\n";
            $state = CRITICAL;
        }
    }

    return ($state,$answer);
}

sub uploadFile {
    my $localurl = shift;
    my $url = shift;
    my $verbose = shift;

    my $state = OK;
    my $answer;
    my $res;

    print "INFO: transferring file to remote computer.\n" if ($verbose);
    print "COMMAND: $COMMANDS{GUC} file:$localurl $url \n" if ($verbose);

    ($state, $res) = processCommand ("$COMMANDS{GUC} file:$localurl $url 2>&1");

    if ($state == CRITICAL) {
        $answer = processErrorOutput($res, "Upload to remote computer failed"). "\n";
    } elsif ($state == UNKNOWN) {
        $answer = $res;
    } else {
        print "INFO: transferring file to remote computer succeeded\n" if ($verbose);
        $answer = "Upload to remote computer succeeded. ";
    }

    return ($state, $answer, $res);
}

sub downloadFile {
    my $localurl = shift;
    my $url = shift;
    my $verbose = shift;

    my $state = OK;
    my $answer;
    my $res;

    print "INFO: transferring file from remote to local computer.\n" if ($verbose);
    print "COMMAND: $COMMANDS{GUC} $url file:${localurl}-res \n" if ($verbose);

    ($state, $res) = processCommand ("$COMMANDS{GUC} $url file:${localurl}-res 2>&1");

    if ($state == CRITICAL) {
        $answer = processErrorOutput($res, "Download from remote computer failed") . "\n";
    } elsif ($state == UNKNOWN) {
        $answer = $res;
    } else {
        print "INFO: transferring file from remote to local computer succeeded. Checking received file...\n" if ($verbose);
        $answer = "Download from remote computer succeeded. ";
    }

    return ($state, $answer, $res);
}

sub deleteFile {
    my $host = shift;
    my $port = shift;
    my $url = shift;
    my $verbose = shift;

    my $state = OK;
    my $answer;
    my $res;

    print "INFO: removing file from remote computer.\n" if ($verbose);
    print "COMMAND: $COMMANDS{UBERFTP} -P $port $host 'rm $url' \n" if ($verbose);

    ($state, $res) = processCommand ("$COMMANDS{UBERFTP} -P $port $host 'rm $url' 2>&1");

    if ($state == CRITICAL) {
        $answer = processErrorOutput($res, "Removing file from remote computer failed") . "\n";
    } elsif ($state == UNKNOWN) {
        $answer = $res;
    } else {
        $answer = "File successfully removed from remote computer. ";
    }

    return ($state, $answer, $res);
}

sub checkFile {
    my $localurl = shift;
    my $data = shift;

    my $state = OK;
    my $answer;

    unless (open (STD, "${localurl}-res")) {
        $answer = "Opening received file failed.\n";
        $state = CRITICAL;
    } else {
        my $newData = <STD>;
        unless (close (STD)) {
            $answer = "Opening received file failed.\n";
            $state = CRITICAL;
        } else {
            if ($newData eq $data) {
                $answer = "Received file is valid. ";
            } else {
                $answer = "Content of received file is not the same as the sent file: \"$newData\".\n";
                $state = CRITICAL;
            }
        }
    }

    return ($state, $answer);
}

sub printOutputClean {
    my $state = shift;
    my $answer = shift;
    my $res = shift;
    my $plugin = shift;
    my $localurl = shift;
    my $url = shift;
    my $host = shift;
    my $port = shift;



    if ($localurl) {
        unlink($localurl);
        unlink("${localurl}-res");
    }

    if ($url) {
        deleteFile($host, $port, $url);
    }

    alarm(0);

    $plugin->nagios_exit($state, $answer);
}

# Local variables
my ($state,$answer,$res,$tmpAnswer);
my ($hostname,$port);
my ($localurl,$serviceurl,$timeoutanswer,$distPath,$localDir);
my $myData = time();

# Create Nagios::Plugin instance
my $plugin = Nagios::Plugin->new (usage => USAGE,
                                  shortname => SHORTNAME,
                                  version => VERSION,
                                  blurb => DESCRIPTION,
                                  extra => EXTRA_DESC,
                                  license => LICENSE,
                                  plugin  => PROGNAME);
# Define additional arguments
$plugin->add_arg(
    spec => 'hostname|H=s',
    help => "H|hostname\n   Name or IP address of host to check.",
    required => 0,
    default => 'localhost'
);
$plugin->add_arg(
    spec => 'port|p=i',
    help => "p|port\n   Port of the service.\n   (default: ".DEFAULT_PORT.")",
    required => 0,
    default => DEFAULT_PORT
);
$plugin->add_arg(
    spec => 'proxy|x=s',
    help => "proxy|x\n   Location of Nagios user's proxy file.\n   (default: /tmp/x509up_u$<)",
    required => 0,
    default => "/tmp/x509up_u$<"
);
$plugin->add_arg(
    spec => 'vo=s',
    help => "vo\n   Virtual organization of user.\n   (default: )",
    required => 0,
);
$plugin->add_arg(
    spec => 'path=s',
    help => "path\n   Remote path to use as destination path.\n   (default: ".DEFAULT_TMP_DIR.")",
    required => 0,
    default => DEFAULT_TMP_DIR
);
$plugin->add_arg(
    spec => 'dir=s',
    help => "dir\n   Local directory used as source path.\n   (default: ".DEFAULT_TMP_DIR.")",
    required => 0,
    default => DEFAULT_TMP_DIR
);

$plugin->getopts;

$hostname = $plugin->opts->hostname;
$port = $plugin->opts->port;

$VONAME = $plugin->opts->vo if ($plugin->opts->vo);

my $globusLocation = $ENV{GLOBUS_LOCATION} || "/usr";
$COMMANDS{GUC}  = "$globusLocation/bin/globus-url-copy";
$COMMANDS{UBERFTP}  = "$globusLocation/bin/uberftp";
#check if defined/required commands are available
#foreach my $key (keys %COMMANDS) {
#	my $file = $COMMANDS{$key};
#	unless (-e $file) {
#		$plugin->nagios_die("Required command $file not available.");
#	}
#}


# Just in case of problems, let's not hang Nagios
local $SIG{'ALRM'} = sub {
    local $SIG{TERM} = 'IGNORE';
    kill TERM => -$$;
    printOutputClean(UNKNOWN, "$timeoutanswer\n","",$plugin,$localurl);
};

local $SIG{TERM} = sub {
    local $SIG{TERM} = 'IGNORE';
    kill TERM => -$$;
    printOutputClean(UNKNOWN, "Plugin received TERM signal.\n","",$plugin,$localurl);
};

alarm($plugin->opts->timeout);

$localDir = $plugin->opts->dir;
`/bin/mkdir -p $localDir` unless (-d $localDir);
unless (-w $localDir) {
    $plugin->nagios_die("GridFTP cannot create local directory $localDir.");
}

($state,$answer,$res) = checkProxy($plugin->opts->proxy);
if ( $state != OK) {
    $plugin->nagios_exit($state, $answer);
}

$serviceurl = 'gsiftp://' . $hostname . ':' . $port . '/' . $plugin->opts->path;
$distPath = $plugin->opts->path;
$distPath .= "/GridFTP-Probe-" . hostname() . "-$hostname-$port-$$";
$serviceurl .= "/GridFTP-Probe-" . hostname() . "-$hostname-$port-$$";
$localurl = "$localDir/GridFTP-Probe-$hostname-$port-$$";

$state = OK;

($state,$answer) = createFile($localurl,$myData);
if ( $state != OK) {
    $plugin->nagios_exit($state,$answer);
}
$timeoutanswer = "Timeout occured during upload to remote computer.";
($state,$answer,$res) = uploadFile($localurl, $serviceurl, $plugin->opts->verbose);
if ( $state != OK) {
    printOutputClean ($state,$answer,$res,$plugin,$localurl);
}
$timeoutanswer = "Timeout occured during download from remote computer.";
($state,$tmpAnswer,$res) = downloadFile($localurl, $serviceurl, $plugin->opts->verbose);
$answer .= $tmpAnswer;
if ( $state != OK) {
    printOutputClean ($state,$answer,$res,$plugin,$localurl,$serviceurl,$hostname,$port);
}
$timeoutanswer = "Timeout occured during download from remote computer.";
($state,$tmpAnswer,$res) = deleteFile($hostname, $port, $distPath, $plugin->opts->verbose);
$answer .= $tmpAnswer;
if ( $state != OK) {
    printOutputClean ($state,$answer,$res,$plugin,$localurl);
}
($state,$tmpAnswer,$res) = checkFile($localurl, $myData);
$answer .= $tmpAnswer;

printOutputClean ($state,$answer,$res,$plugin,$localurl);

